import argparse
import os
import numpy as np
import tensorflow as tf
from matplotlib import pyplot as plt
from PIL import Image
import datetime
from utils import visualization_utils as vis_util
from utils import label_map_util
from xml_factory import detection2xml
if tf.__version__ < '1.4.0':
    raise ImportError('Please upgrade your tensorflow installation to v1.4.* or later!')

NUM_CLASSES = 90

'''使用说明：
1.output_dir：检测后的输出图片保存路径
2.model_dir：使用的检测模型计算图保存路径
3.input_dir：输入图片路径
4.fail_dir：未给成功检测的图片输出路径
5.xml_dir：输出xml的保存路径
6.labels:labels.txt文件路径
例：python inference_relabel.py --output_dir=out --model_dir=model --input_dir=input --fail_dir=fail --xml_dir=xmls --labels=labels.txt'''
def parse_args(check=True):
    parser = argparse.ArgumentParser()
    parser.add_argument('--output_dir', type=str, required=True)
    parser.add_argument('--model_dir', type=str, required=True)
    parser.add_argument('--input_dir', type=str, required=True)
    parser.add_argument('--fail_dir', type=str, required=True)
    parser.add_argument('--xml_dir', type=str, required=True)
    parser.add_argument('--labels', type=str, required=True)
    FLAGS, unparsed = parser.parse_known_args()
    return FLAGS, unparsed

def check_result(boxes, scores, classes):
    i = 0
    for score in scores:
        if score > 0.5:
            i=i+1
        else:
            break
    if i == 0:
        return -1
    scores = scores[0:i]
    indexList=[]
    areaList=[]
    for index in range(i):
        if int(classes[index]) == 3 or int(classes[index]) == 8:
            indexList.append(index)
            areaList.append((boxes[index][2]-boxes[index][0])*(boxes[index][3]-boxes[index][1]))
    if indexList:
        return indexList[areaList.index(max(areaList))]
    else:
        return -1


#print(labeldic)
if __name__ == '__main__':
    FLAGS, unparsed = parse_args()

    PATH_TO_CKPT = os.path.join(FLAGS.model_dir, 'frozen_inference_graph.pb')
    PATH_TO_LABELS = os.path.join(FLAGS.model_dir, 'mscoco_label_map.pbtxt')
    if not os.path.exists(FLAGS.output_dir):
        os.makedirs(FLAGS.output_dir)

    if not os.path.exists(FLAGS.fail_dir):
        os.makedirs(FLAGS.fail_dir)

    if not os.path.exists(FLAGS.xml_dir):
        os.makedirs(FLAGS.xml_dir)
    labeldic = {}
    with open(FLAGS.labels, 'rb') as f:
        for line in f:
            line = line.decode('utf-8').strip()
            labeldic[line.split(':')[0]] = line.split(':')[1]
    detection_graph = tf.Graph()
    with detection_graph.as_default():
        od_graph_def = tf.GraphDef()
        with tf.gfile.GFile(PATH_TO_CKPT, 'rb') as fid:
            serialized_graph = fid.read()
            od_graph_def.ParseFromString(serialized_graph)
            tf.import_graph_def(od_graph_def, name='')

    label_map = label_map_util.load_labelmap(PATH_TO_LABELS)
    categories = label_map_util.convert_label_map_to_categories(label_map, max_num_classes=NUM_CLASSES, use_display_name=True)
    category_index = label_map_util.create_category_index(categories)

    def load_image_into_numpy_array(image):
        (im_width, im_height) = image.size
        return np.array(image.getdata()).reshape(
            (im_height, im_width, 3)).astype(np.uint8)
    
    for root, dirs, files in os.walk(FLAGS.input_dir):
        pass
    #print(files)
    count = 0
    starttime = datetime.datetime.now()
    for mfile in files:
        test_img_path = os.path.join(FLAGS.input_dir, mfile)
        (filename,extension) = os.path.splitext(mfile)
        xmlpath = os.path.join(FLAGS.xml_dir, filename + '.xml')
        label = filename.split('_')[2]
        with detection_graph.as_default():
            with tf.Session(graph=detection_graph) as sess:
                coord=tf.train.Coordinator()
                threads= tf.train.start_queue_runners(coord=coord)
                image_tensor = detection_graph.get_tensor_by_name('image_tensor:0')
                detection_boxes = detection_graph.get_tensor_by_name('detection_boxes:0')
                detection_scores = detection_graph.get_tensor_by_name('detection_scores:0')
                detection_classes = detection_graph.get_tensor_by_name('detection_classes:0')
                num_detections = detection_graph.get_tensor_by_name('num_detections:0')
                image = Image.open(test_img_path)
                image_np = load_image_into_numpy_array(image)
                image_np_expanded = np.expand_dims(image_np, axis=0)
                (boxes, scores, classes, num) = sess.run(
                    [detection_boxes, detection_scores, detection_classes, num_detections],
                    feed_dict={image_tensor: image_np_expanded})
                boxes = np.squeeze(boxes)
                classes = np.squeeze(classes)
                scores = np.squeeze(scores)
                index = check_result(boxes, scores, classes)
                if index > -1:
                    xmin = boxes[index:index + 1][0][1] * image.size[0]
                    xmax = boxes[index:index + 1][0][3] * image.size[0]
                    ymin = boxes[index:index + 1][0][0] * image.size[1]
                    ymax = boxes[index:index + 1][0][2] * image.size[1]
                    detection2xml(xmlpath,[image.size[0],image.size[1],int(3)],[[labeldic[label],int(xmin),int(ymin),int(xmax),int(ymax)],])
                    print("==%d==index=%d==class=%d=====score=%f=====%s" % (count,index,int(classes[index:index+1]),scores[index:index+1],filename))
                    #print(classes[index:index+1])
                    #print(scores[index:index+1])
                    vis_util.visualize_boxes_and_labels_on_image_array(
                        image_np,
                        boxes[index:index+1],
                        classes[index:index+1].astype(np.int32),
                        scores[index:index+1],
                        category_index,
                        use_normalized_coordinates=True,
                        line_thickness=8)
                    #print(filename)
                    outpath = os.path.join(FLAGS.output_dir, filename+'.png')
                    #print(outpath)
                    plt.imsave(outpath, image_np)
                else:
                    cmd = 'cp '+test_img_path+' '+FLAGS.fail_dir
                    print(cmd)
                    os.system(cmd)
                coord.request_stop()
                coord.join(threads)
        count=count+1
        if count % 10 == 0:
            nowtime = datetime.datetime.now()
            print("=%d image==cost %s seconds===%fs/p" % ((count,(nowtime-starttime).seconds,(nowtime-starttime).seconds/count)))
